import chisel3._
import chiseltest._
import org.scalatest.{Matchers, FlatSpec}
import scala.util.Random

import comlib._

class CrossbarSpec extends FlatSpec with ChiselScalatestTester with Matchers {

  class MyModule(inFF:Boolean, selFF:Int) extends Module{
    val io = IO(new Bundle{
      val in =  Input(Vec(4, UInt(16.W)))
      val sel = Input(Vec(4, UInt(2.W)))
      val out = Output(Vec(4, UInt(16.W)))
    })

    val xb = Crossbar.xbB( io.in, io.sel, inFF=inFF, selFF )
    io.out := xb.out 
  }

  val rnd = new Random
  val sel = Seq(0, 1, 2, 3)

  it should "support no FF application" in {
    test( new MyModule(false, 0) ){ dut => 
      var idata = 0

      for( it <- 0 until 100 ){
        val sels = rnd.shuffle(sel)

        dut.io.in.foreach { case i => 
          i.poke(idata.U)
          idata = idata + 1
        }

        dut.io.sel.zip(sels).foreach { case (io, sel) => io.poke(sel.U) }

        dut.io.out.zip(sels).foreach { case (out,sel) =>
          out.expect((idata-4+sel).U, f"sel=${sel} sels=${sels} it=${it}")
        }
      }

    }
  }

  it should "support 1-FFs for in and sel signals application" in {
    test( new MyModule(true, 1) ){ dut => 
      var idata = 0

      for( it <- 0 until 100 ){
        val sels = rnd.shuffle(sel)

        dut.io.in.foreach { case i => 
          i.poke(idata.U)
          idata = idata + 1
        }

        dut.io.sel.zip(sels).foreach { case (io, sel) => io.poke(sel.U) }

        dut.clock.step()
        dut.io.out.zip(sels).foreach { case (out,sel) =>
          out.expect((idata-4+sel).U, f"sel=${sel} sels=${sels} it=${it}")
        }
      }

    }
  }

  it should "support average split crossbar application" in {
    test( new MyModule(true, 4) ){ dut => 
      var idata = 0

      for( it <- 0 until 100 ){
        val sels = rnd.shuffle(sel)

        dut.io.in.foreach { case i => 
          i.poke(idata.U)
          idata = idata + 1
        }

        dut.io.sel.zip(sels).foreach { case (io, sel) => io.poke(sel.U) }

        dut.clock.step()
        dut.io.out.zip(sels).foreach { case (out,sel) =>
          out.expect((idata-4+sel).U, f"sel=${sel} sels=${sels} it=${it}")
        }
      }

    }
  }

  it should "support non-average split crossbar application" in {
    test( new MyModule(true, 3) ){ dut => 
      var idata = 0

      for( it <- 0 until 100 ){
        val sels = rnd.shuffle(sel)

        dut.io.in.foreach { case i => 
          i.poke(idata.U)
          idata = idata + 1
        }

        dut.io.sel.zip(sels).foreach { case (io, sel) => io.poke(sel.U) }

        dut.clock.step()
        dut.io.out.zip(sels).foreach { case (out,sel) =>
          out.expect((idata-4+sel).U, f"sel=${sel} sels=${sels} it=${it}")
        }
      }

    }
  }

  class MyModule2 extends Module{

    class MyPacket extends Bundle{
      val vld = Bool()
      val data = UInt(16.W)
      val sn = UInt(4.W)
    }
    object MyPacket{
      def apply(data:UInt) = {
        val res = Wire(new MyPacket)
        res.vld := data(0)
        res.data := data 
        res.sn := data(3,0)
        res
      }
    }

    val io = IO(new Bundle{
      val in =  Input(Vec(4, UInt(16.W)))
      val sel = Input(Vec(4, UInt(2.W)))
      val out = Output(Vec(4, UInt(16.W)))
    })

    val in_inner = io.in.map{ case in => MyPacket(in) }
    val xb = Crossbar.xbB( VecInit(in_inner), io.sel, inFF=true, selFF=3 )
    io.out := VecInit( xb.out.map{ case o => o.data } )
  }

  it should "support Bundle non-average split crossbar application" in {
    test( new MyModule2 ){ dut => 
      var idata = 0

      for( it <- 0 until 100 ){
        val sels = rnd.shuffle(sel)

        dut.io.in.foreach { case i => 
          i.poke(idata.U)
          idata = idata + 1
        }

        dut.io.sel.zip(sels).foreach { case (io, sel) => io.poke(sel.U) }

        dut.clock.step()
        dut.io.out.zip(sels).foreach { case (out,sel) =>
          out.expect((idata-4+sel).U, f"sel=${sel} sels=${sels} it=${it}")
        }
      }

    }
  }

}